"""
人脸质量检测模块 - 改进版本
"""
import cv2
import numpy as np
from typing import Tuple, Dict, Optional
from pathlib import Path
from loguru import logger
from model_manager import model_manager
from config import settings


class FaceQualityDetector:
    """人脸质量检测器"""
    
    def __init__(self):
        """初始化检测器"""
        self.model = model_manager.get_model()
    
    def detect_from_path(
        self, 
        image_path: str,
        min_face_size: Optional[int] = None,
        max_angle: Optional[float] = None,
        blur_threshold: Optional[float] = None,
        min_quality_score: Optional[float] = None
    ) -> Tuple[bool, str, Optional[Dict]]:
        """
        从文件路径检测人脸质量
        
        Args:
            image_path: 图片路径
            min_face_size: 最小人脸尺寸（像素）
            max_angle: 最大允许偏转角度
            blur_threshold: 模糊阈值
            min_quality_score: 最低质量评分阈值（0-1）
            
        Returns:
            (是否合格, 消息, 详细信息)
        """
        # 加载图像
        img = cv2.imread(image_path)
        if img is None:
            return False, "无法加载图片", None
        
        return self.detect_from_array(img, min_face_size, max_angle, blur_threshold, min_quality_score)
    
    def detect_from_bytes(
        self,
        image_bytes: bytes,
        min_face_size: Optional[int] = None,
        max_angle: Optional[float] = None,
        blur_threshold: Optional[float] = None,
        min_quality_score: Optional[float] = None
    ) -> Tuple[bool, str, Optional[Dict]]:
        """
        从字节数据检测人脸质量
        
        Args:
            image_bytes: 图片字节数据
            min_face_size: 最小人脸尺寸
            max_angle: 最大允许偏转角度
            blur_threshold: 模糊阈值
            min_quality_score: 最低质量评分阈值（0-1）
            
        Returns:
            (是否合格, 消息, 详细信息)
        """
        # 将字节转换为numpy数组
        nparr = np.frombuffer(image_bytes, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img is None:
            return False, "无法解析图片数据", None
        
        return self.detect_from_array(img, min_face_size, max_angle, blur_threshold, min_quality_score)
    
    def detect_from_array(
        self,
        img: np.ndarray,
        min_face_size: Optional[int] = None,
        max_angle: Optional[float] = None,
        blur_threshold: Optional[float] = None,
        min_quality_score: Optional[float] = None
    ) -> Tuple[bool, str, Optional[Dict]]:
        """
        从numpy数组检测人脸质量
        
        Args:
            img: 图片numpy数组
            min_face_size: 最小人脸尺寸
            max_angle: 最大允许偏转角度  
            blur_threshold: 模糊阈值
            min_quality_score: 最低质量评分阈值（0-1）
            
        Returns:
            (是否合格, 消息, 详细信息)
        """
        # 使用配置中的默认值
        min_face_size = min_face_size or settings.min_face_size
        max_angle = max_angle or settings.max_angle
        blur_threshold = blur_threshold or settings.blur_threshold
        min_quality_score = min_quality_score if min_quality_score is not None else settings.min_quality_score
        
        try:
            # 检测人脸
            faces = self.model.get(img)
            
            if len(faces) == 0:
                return False, "未检测到人脸", None
            
            if len(faces) > 1:
                return False, f"检测到多张人脸 ({len(faces)}张)", None
            
            face = faces[0]
            
            # 提取人脸区域
            x1, y1, x2, y2 = [int(v) for v in face.bbox]
            w, h = x2 - x1, y2 - y1
            
            # 人脸尺寸检查
            if w < min_face_size or h < min_face_size:
                return False, f"人脸太小 ({w}x{h}像素)", {
                    "width": int(w),
                    "height": int(h),
                    "min_required": int(min_face_size)
                }
            
            # 清晰度检查
            gray = cv2.cvtColor(img[y1:y2, x1:x2], cv2.COLOR_BGR2GRAY)
            blur_score = cv2.Laplacian(gray, cv2.CV_64F).var()
            
            if blur_score < blur_threshold:
                return False, f"人脸模糊 (清晰度={blur_score:.2f})", {
                    "blur_score": float(blur_score),
                    "threshold": float(blur_threshold)
                }
            
            # 姿态角度检查
            yaw, pitch, roll = face.pose
            
            if abs(yaw) > max_angle or abs(pitch) > max_angle or abs(roll) > max_angle:
                return False, f"人脸角度过大", {
                    "yaw": float(yaw),
                    "pitch": float(pitch),
                    "roll": float(roll),
                    "max_angle": float(max_angle)
                }
            
            # 质量评分
            quality_score = self._calculate_quality_score(
                blur_score, blur_threshold,
                abs(yaw), abs(pitch), abs(roll), max_angle,
                w, h, min_face_size
            )
            
            # 质量评分检查
            if quality_score < min_quality_score:
                return False, f"质量评分过低 (评分={quality_score:.3f})", {
                    "quality_score": float(quality_score),
                    "min_required": float(min_quality_score),
                    "blur_score": float(blur_score),
                    "face_size": {"width": int(w), "height": int(h)},
                    "pose": {
                        "yaw": float(yaw),
                        "pitch": float(pitch),
                        "roll": float(roll)
                    }
                }
            
            # 所有检查通过
            details = {
                "face_size": {"width": int(w), "height": int(h)},
                "blur_score": float(blur_score),
                "pose": {
                    "yaw": float(yaw),
                    "pitch": float(pitch),
                    "roll": float(roll)
                },
                "quality_score": float(quality_score),
                "bbox": [int(x1), int(y1), int(x2), int(y2)],
                "landmarks": face.kps.tolist() if hasattr(face, 'kps') else None,
                "age": int(face.age) if hasattr(face, 'age') else None,
                "gender": int(face.gender) if hasattr(face, 'gender') else None
            }
            
            return True, "标准人脸 ✅", details
            
        except Exception as e:
            logger.error(f"人脸检测失败: {e}")
            return False, f"检测过程出错: {str(e)}", None
    
    def _calculate_quality_score(
        self,
        blur_score: float, blur_threshold: float,
        yaw: float, pitch: float, roll: float, max_angle: float,
        width: int, height: int, min_face_size: int
    ) -> float:
        """
        计算综合质量评分 (0-1)
        
        Args:
            blur_score: 清晰度分数
            blur_threshold: 清晰度阈值
            yaw, pitch, roll: 姿态角度
            max_angle: 最大角度
            width, height: 人脸尺寸
            min_face_size: 最小尺寸要求
            
        Returns:
            质量评分 (0-1)
        """
        # 清晰度得分 (0-1)
        blur_score_norm = min(blur_score / (blur_threshold * 3), 1.0)
        
        # 姿态得分 (0-1)
        avg_angle = (abs(yaw) + abs(pitch) + abs(roll)) / 3
        pose_score = max(0, 1 - (avg_angle / max_angle))
        
        # 尺寸得分 (0-1)
        avg_size = (width + height) / 2
        size_score = min(avg_size / (min_face_size * 3), 1.0)
        
        # 综合评分 (加权平均)
        quality = (blur_score_norm * 0.4 + pose_score * 0.4 + size_score * 0.2)
        
        return round(quality, 3)


# 创建全局检测器实例
face_detector = FaceQualityDetector()

